package main

import (
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/alecthomas/kingpin"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	log "github.com/sirupsen/logrus"

	"github.com/prometheus/common/version"
)

const (
	namespace = "port" // For Prometheus metrics.
)

var (
	listeningAddress = kingpin.Flag("listen.address", "Address on which to expose metrics.").Default(":9120").String()
	metricsEndpoint  = kingpin.Flag("metrics.endpoint", "Path under which to expose metrics.").Default("/metrics").String()
	ports            = kingpin.Flag("ports", "Ports to scrape metrics of connections.Seperated by comma.").Default("").String()
	collectType      = kingpin.Flag("collect.type", "Ports type to collect.").Default("all").String()
	gracefulStop     = make(chan os.Signal)
)

type Exporter struct {
	mutex sync.Mutex
	ports []string

	up               *prometheus.GaugeVec
	connection_total *prometheus.GaugeVec
	connection_count *prometheus.GaugeVec
}

func reverse(s []string) []string {
	for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 {
		s[i], s[j] = s[j], s[i]
	}
	return s
}

func (e *Exporter) get_port_metrics(port string) error {

	// up
	statement := fmt.Sprintf("ss -atnl | grep \":%s[^0-9]\" | wc -l", port)
	log.Debugf("running bash command: %v", statement)
	output, err := exec.Command("sh", "-c", statement).CombinedOutput()
	if err != nil {
		return err
	}
	val, err := strconv.ParseFloat(strings.TrimSpace(string(output)), 64)
	if err != nil {
		return err
	}
	e.up.WithLabelValues(port).Set(val)

	if *collectType != "port_up" {
		// establish / timewait / closewait
		port_states := []string{"ESTAB", "TIME-WAIT", "CLOSE-WAIT"}
		for _, state := range port_states {
			statement := fmt.Sprintf("ss -atn | grep \":%s[^0-9]\" | grep %s", port, state)
			log.Debugf("running bash command: %v", statement)

			lines := []string{}
			output, err := exec.Command("sh", "-c", statement).CombinedOutput()
			if err != nil && err.Error() != "exit status 1" {
				return err
			}
			lines = strings.Split(string(output), "\n")
			e.connection_total.WithLabelValues(port, state).Set(float64(len(lines)) - 1)
			dest_map := make(map[string]float64)
			for _, line := range lines {
				re, _ := regexp.Compile(`\s+`)
				ss := re.Split(line, -1)
				if len(ss) > 1 {
					dest_arr := strings.Split(string(ss[4]), ":")
					dest_arr = reverse(dest_arr)
					dest := dest_arr[1]
					if _, ok := dest_map[dest]; !ok { // dest not exists in dest_map
						dest_map[dest] = 0
					}
					dest_map[dest] = dest_map[dest] + 1
				}
			}
			for key, value := range dest_map {
				e.connection_count.WithLabelValues(port, key, state).Set(value)
			}
		}
	}

	return nil
}
func NewExporter(ports []string) *Exporter {
	if *collectType != "port_up" {
		return &Exporter{
			ports: ports,
			up: prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "up",
				Help:      "If the port is up", // metrics help info
			},
				[]string{"port"},
			),

			connection_total: prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "connection_total",
				Help:      "The total count of port connection", // metrics help info
			},
				[]string{"port", "state"},
			),
			connection_count: prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "connection_count",
				Help:      "The count of port connection with destinatins", // metrics help info
			},
				[]string{"port", "dest", "state"},
			),
		}
	} else {
		return &Exporter{
			ports: ports,
			up: prometheus.NewGaugeVec(prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "up",
				Help:      "If the port is up", // metrics help info
			},
				[]string{"port"},
			),
		}
	}

}

func (e *Exporter) scrape_metrics(ch chan<- prometheus.Metric) error {
	if e.ports[0] == "" {
		statement := "ss -atnl | grep -v State | awk '{print $4}'"
		log.Debugf("running bash command: %v", statement)
		output, err := exec.Command("sh", "-c", statement).CombinedOutput()
		if err != nil {
			os.Exit(-1)
		}
		lines := strings.Split(string(output), "\n")
		for _, line := range lines {
			if line != "" {
				arr := strings.Split(line, ":")
				port := arr[len(arr)-1]
				if err = e.get_port_metrics(port); err != nil {
					return fmt.Errorf("Error scrape port metrics: %v", err)
				}
			}
		}
	} else {
		for _, port := range e.ports {
			if err := e.get_port_metrics(port); err != nil {
				return fmt.Errorf("Error scrape port metrics: %v", err)
			}
		}
	}
	e.up.Collect(ch)
	if *collectType != "port_up" {
		e.connection_total.Collect(ch)
		e.connection_count.Collect(ch)
	}

	return nil
}

func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()
	if err := e.scrape_metrics(ch); err != nil {
		log.Errorf("%v", err)
	}
	return
}

func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	e.up.Describe(ch)
	if *collectType != "port_up" {
		e.connection_total.Describe(ch)
		e.connection_count.Describe(ch)
	}
}

func main() {
	// Parse flags
	log.SetLevel(log.InfoLevel)
	kingpin.Version(version.Print("port_exporter"))
	kingpin.HelpFlag.Short('h')
	kingpin.Parse()

	// listen to termination signals from the OS
	signal.Notify(gracefulStop, syscall.SIGTERM)
	signal.Notify(gracefulStop, syscall.SIGINT)
	signal.Notify(gracefulStop, syscall.SIGHUP)
	signal.Notify(gracefulStop, syscall.SIGQUIT)

	log.Infoln("Starting port_exporter", version.Info())
	log.Infoln("Build context", version.BuildContext())
	log.Infof("Starting Server: %s", *listeningAddress)
	portArr := strings.Split(*ports, ",")
	log.Infof("Start to get metrics of ports: %v", portArr)

	exporter := NewExporter(portArr)
	re := prometheus.NewRegistry()
	re.MustRegister(exporter)
	// re.MustRegister(version.NewCollector("port_exporter"))

	// listener for the termination signals from the OS
	go func() {
		log.Infof("listening and wait for graceful stop")
		sig := <-gracefulStop
		log.Infof("caught sig: %+v. Wait 1 seconds...", sig)
		time.Sleep(1 * time.Second)
		log.Infof("Terminate port_exporter on port: %s", *listeningAddress)
		os.Exit(0)
	}()

	handler := promhttp.HandlerFor(re, promhttp.HandlerOpts{})
	http.Handle(*metricsEndpoint, handler)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte(`<html>
		<head><title>Port Exporter</title></head>
		<body>
		<h1>Port Exporter</h1>
		<p>Metrics</p >
		</body>
		</html>`))
	})
	log.Fatal(http.ListenAndServe(*listeningAddress, nil))
}
